// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <cstdint>
#include <memory>
#include <vector>

#include <nebula/core/array.h>
#include <nebula/core/record_batch.h>
#include <turbo/utility/status.h>
#include <nebula/types/type.h>

namespace nebula {
namespace ipc {
namespace test {

// A typedef used for test parameterization
typedef turbo::Status MakeRecordBatch(std::shared_ptr<RecordBatch>* out);

TURBO_EXPORT
void CompareArraysDetailed(int index, const Array& result, const Array& expected);

TURBO_EXPORT
void CompareBatchColumnsDetailed(const RecordBatch& result, const RecordBatch& expected);

TURBO_EXPORT
turbo::Status MakeRandomInt32Array(int64_t length, bool include_nulls, MemoryPool* pool,
                            std::shared_ptr<Array>* out, uint32_t seed = 0,
                            int32_t min = 0, int32_t max = 1000);

TURBO_EXPORT
turbo::Status MakeRandomInt64Array(int64_t length, bool include_nulls, MemoryPool* pool,
                            std::shared_ptr<Array>* out, uint32_t seed = 0);

TURBO_EXPORT
turbo::Status MakeRandomListArray(const std::shared_ptr<Array>& child_array, int num_lists,
                           bool include_nulls, MemoryPool* pool,
                           std::shared_ptr<Array>* out);

TURBO_EXPORT
turbo::Status MakeRandomLargeListArray(const std::shared_ptr<Array>& child_array, int num_lists,
                                bool include_nulls, MemoryPool* pool,
                                std::shared_ptr<Array>* out);

TURBO_EXPORT
turbo::Status MakeRandomBooleanArray(const int length, bool include_nulls,
                              std::shared_ptr<Array>* out);

TURBO_EXPORT
turbo::Status MakeBooleanBatchSized(const int length, std::shared_ptr<RecordBatch>* out);

TURBO_EXPORT
turbo::Status MakeBooleanBatch(std::shared_ptr<RecordBatch>* out);

TURBO_EXPORT
turbo::Status MakeIntBatchSized(int length, std::shared_ptr<RecordBatch>* out,
                         uint32_t seed = 0);

TURBO_EXPORT
turbo::Status MakeIntRecordBatch(std::shared_ptr<RecordBatch>* out);

TURBO_EXPORT
turbo::Status MakeFloat3264BatchSized(int length, std::shared_ptr<RecordBatch>* out,
                               uint32_t seed = 0);

TURBO_EXPORT
turbo::Status MakeFloat3264Batch(std::shared_ptr<RecordBatch>* out);

TURBO_EXPORT
turbo::Status MakeFloatBatchSized(int length, std::shared_ptr<RecordBatch>* out,
                           uint32_t seed = 0);

TURBO_EXPORT
turbo::Status MakeFloatBatch(std::shared_ptr<RecordBatch>* out);

TURBO_EXPORT
turbo::Status MakeRandomStringArray(int64_t length, bool include_nulls, MemoryPool* pool,
                             std::shared_ptr<Array>* out);

TURBO_EXPORT
turbo::Status MakeStringTypesRecordBatch(std::shared_ptr<RecordBatch>* out,
                                  bool with_nulls = true, bool with_view_types = true);

TURBO_EXPORT
turbo::Status MakeStringTypesRecordBatchWithNulls(std::shared_ptr<RecordBatch>* out);

TURBO_EXPORT
turbo::Status MakeNullRecordBatch(std::shared_ptr<RecordBatch>* out);

TURBO_EXPORT
turbo::Status MakeListRecordBatch(std::shared_ptr<RecordBatch>* out);

TURBO_EXPORT
turbo::Status MakeListViewRecordBatch(std::shared_ptr<RecordBatch>* out);

TURBO_EXPORT
turbo::Status MakeFixedSizeListRecordBatch(std::shared_ptr<RecordBatch>* out);

TURBO_EXPORT
turbo::Status MakeZeroLengthRecordBatch(std::shared_ptr<RecordBatch>* out);

TURBO_EXPORT
turbo::Status MakeNonNullRecordBatch(std::shared_ptr<RecordBatch>* out);

TURBO_EXPORT
turbo::Status MakeDeeplyNestedList(std::shared_ptr<RecordBatch>* out);

TURBO_EXPORT
turbo::Status MakeDeeplyNestedListView(std::shared_ptr<RecordBatch>* out);

TURBO_EXPORT
turbo::Status MakeStruct(std::shared_ptr<RecordBatch>* out);

TURBO_EXPORT
turbo::Status MakeRunEndEncoded(std::shared_ptr<RecordBatch>* out);

TURBO_EXPORT
turbo::Status MakeUnion(std::shared_ptr<RecordBatch>* out);

TURBO_EXPORT
turbo::Status MakeDictionary(std::shared_ptr<RecordBatch>* out);

TURBO_EXPORT
turbo::Status MakeDictionaryFlat(std::shared_ptr<RecordBatch>* out);

TURBO_EXPORT
turbo::Status MakeNestedDictionary(std::shared_ptr<RecordBatch>* out);

TURBO_EXPORT
turbo::Status MakeMap(std::shared_ptr<RecordBatch>* out);

TURBO_EXPORT
turbo::Status MakeMapOfDictionary(std::shared_ptr<RecordBatch>* out);

TURBO_EXPORT
turbo::Status MakeDates(std::shared_ptr<RecordBatch>* out);

TURBO_EXPORT
turbo::Status MakeTimestamps(std::shared_ptr<RecordBatch>* out);

TURBO_EXPORT
turbo::Status MakeIntervals(std::shared_ptr<RecordBatch>* out);

TURBO_EXPORT
turbo::Status MakeTimes(std::shared_ptr<RecordBatch>* out);

TURBO_EXPORT
turbo::Status MakeFWBinary(std::shared_ptr<RecordBatch>* out);

TURBO_EXPORT
turbo::Status MakeDecimal(std::shared_ptr<RecordBatch>* out);

TURBO_EXPORT
turbo::Status MakeNull(std::shared_ptr<RecordBatch>* out);

TURBO_EXPORT
turbo::Status MakeUuid(std::shared_ptr<RecordBatch>* out);

TURBO_EXPORT
turbo::Status MakeComplex128(std::shared_ptr<RecordBatch>* out);

TURBO_EXPORT
turbo::Status MakeDictExtension(std::shared_ptr<RecordBatch>* out);

TURBO_EXPORT
turbo::Status MakeRandomTensor(const std::shared_ptr<DataType>& type,
                        const std::vector<int64_t>& shape, bool row_major_p,
                        std::shared_ptr<Tensor>* out, uint32_t seed = 0);

}  // namespace test
}  // namespace ipc
}  // namespace nebula
